import os
from launch import LaunchDescription
from launch.actions import DeclareLaunchArgument, IncludeLaunchDescription, GroupAction
from launch.substitutions import LaunchConfiguration, PathJoinSubstitution
from launch_ros.actions import Node
from launch_ros.substitutions import FindPackageShare
from launch.launch_description_sources import PythonLaunchDescriptionSource
from ament_index_python.packages import get_package_share_directory

def generate_launch_description():
    """
    Complete 3D navigation launch file integrating ego-planner with existing system
    """
    
    # Package directories
    ego_adapter_dir = get_package_share_directory('ego_planner_adapter')
    ego_planner_dir = get_package_share_directory('ego_planner')
    pb_nav_dir = get_package_share_directory('pb2025_nav_bringup')
    
    # Launch arguments
    namespace_arg = DeclareLaunchArgument(
        'namespace',
        default_value='',
        description='Robot namespace'
    )
    
    use_sim_time_arg = DeclareLaunchArgument(
        'use_sim_time',
        default_value='false',
        description='Use simulation time'
    )
    
    drone_id_arg = DeclareLaunchArgument(
        'drone_id',
        default_value='0',
        description='Drone ID for multi-robot systems'
    )
    
    map_size_x_arg = DeclareLaunchArgument(
        'map_size_x',
        default_value='20.0',
        description='Map size X dimension'
    )
    
    map_size_y_arg = DeclareLaunchArgument(
        'map_size_y',
        default_value='20.0',
        description='Map size Y dimension'
    )
    
    map_size_z_arg = DeclareLaunchArgument(
        'map_size_z',
        default_value='5.0',
        description='Map size Z dimension'
    )
    
    # Parameters
    namespace = LaunchConfiguration('namespace')
    use_sim_time = LaunchConfiguration('use_sim_time')
    drone_id = LaunchConfiguration('drone_id')
    map_size_x = LaunchConfiguration('map_size_x')
    map_size_y = LaunchConfiguration('map_size_y')
    map_size_z = LaunchConfiguration('map_size_z')
    
    # 1. Basic sensor and localization stack (reuse existing)
    basic_navigation = IncludeLaunchDescription(
        PythonLaunchDescriptionSource([
            PathJoinSubstitution([
                FindPackageShare('pb2025_nav_bringup'),
                'launch',
                'localization_launch.py'
            ])
        ]),
        launch_arguments={
            'use_sim_time': use_sim_time,
            'namespace': namespace
        }.items()
    )
    
    # 2. Ego-planner main node
    ego_planner_node = Node(
        package='ego_planner',
        executable='ego_planner_node',
        name='ego_planner',
        namespace=namespace,
        parameters=[
            {'use_sim_time': use_sim_time},
            # Planning parameters
            {'fsm/flight_type': 1},  # 1 for manual waypoint, 2 for preset
            {'fsm/thresh_replan_time': 1.0},
            {'fsm/thresh_no_replan_meter': 1.0},
            {'fsm/planning_horizon': 7.5},
            {'fsm/planning_horizen_time': 3.0},
            {'fsm/emergency_time': 1.0},
            {'fsm/realworld_experiment': True},
            {'fsm/fail_safe': True},
            
            # Map parameters
            {'grid_map/resolution': 0.1},
            {'grid_map/map_size_x': map_size_x},
            {'grid_map/map_size_y': map_size_y},
            {'grid_map/map_size_z': map_size_z},
            {'grid_map/local_update_range_x': 5.5},
            {'grid_map/local_update_range_y': 5.5},
            {'grid_map/local_update_range_z': 4.5},
            {'grid_map/obstacles_inflation': 0.099},
            {'grid_map/local_map_margin': 50},
            {'grid_map/ground_height': -1.0},
            
            # Camera parameters (using LiDAR)
            {'grid_map/cx': 321.04638671875},
            {'grid_map/cy': 243.44969177246094},
            {'grid_map/fx': 387.229248046875},
            {'grid_map/fy': 387.229248046875},
            {'grid_map/use_depth_filter': True},
            {'grid_map/depth_filter_tolerance': 0.15},
            {'grid_map/depth_filter_maxdist': 5.0},
            {'grid_map/depth_filter_mindist': 0.2},
            {'grid_map/depth_filter_margin': 2},
            {'grid_map/k_depth_scaling_factor': 1000.0},
            {'grid_map/skip_pixel': 2},
            
            # Planning algorithm parameters
            {'manager/max_vel': 2.0},
            {'manager/max_acc': 2.0},
            {'manager/max_jerk': 4.0},
            {'manager/dynamic_environment': 0},
            {'manager/local_segment_length': 6.0},
            {'manager/clearance_threshold': 0.2},
            {'manager/control_points_distance': 0.4},
            
            # Optimization parameters
            {'optimization/lambda_smooth': 1.0},
            {'optimization/lambda_collision': 0.5},
            {'optimization/lambda_feasibility': 0.1},
            {'optimization/lambda_fitness': 1.0},
            {'optimization/dist0': 0.5},
            {'optimization/swarm_clearance': 0.5},
            {'optimization/max_vel': 2.0},
            {'optimization/max_acc': 2.0}
        ],
        remappings=[
            ('odom_world', 'odometry'),
            ('planning/bspline', 'ego_planner/trajectory'),
            ('planning/data_display', 'ego_planner/data_display'),
            ('planning/broadcast_bspline_from_planner', 'ego_planner/broadcast_bspline'),
            ('planning/broadcast_bspline_to_planner', 'ego_planner/swarm_trajs'),
            ('grid_map/odom', 'odometry'),
            ('grid_map/cloud', 'cloud_registered'),
            ('grid_map/pose', 'ego_planner/goal'),
            ('grid_map/depth', 'camera/depth/image_rect_raw'),  # Not used with LiDAR
            ('planning/waypoint', 'ego_planner/goal'),
            ('planning/trigger', 'ego_planner/trigger')
        ],
        output='screen'
    )
    
    # 3. Ego-planner adapter
    ego_adapter = IncludeLaunchDescription(
        PythonLaunchDescriptionSource([
            PathJoinSubstitution([
                FindPackageShare('ego_planner_adapter'),
                'launch',
                'ego_planner_adapter.launch.py'
            ])
        ]),
        launch_arguments={
            'namespace': namespace,
            'use_sim_time': use_sim_time
        }.items()
    )
    
    # 4. Serial interface (reuse existing)
    serial_interface = IncludeLaunchDescription(
        PythonLaunchDescriptionSource([
            PathJoinSubstitution([
                FindPackageShare('pb2025_nav_bringup'),
                'launch',
                'serial_interface_launch.py'
            ])
        ]),
        launch_arguments={
            'use_sim_time': use_sim_time,
            'namespace': namespace
        }.items()
    )
    
    # 5. RViz with custom config for 3D navigation
    rviz_config = PathJoinSubstitution([
        FindPackageShare('ego_planner_adapter'),
        'rviz',
        'ego_planner_3d_nav.rviz'
    ])
    
    rviz_node = Node(
        package='rviz2',
        executable='rviz2',
        name='rviz2',
        arguments=['-d', rviz_config],
        parameters=[{'use_sim_time': use_sim_time}],
        output='screen'
    )
    
    # Group everything for clean organization
    ego_3d_navigation = GroupAction([
        basic_navigation,
        ego_planner_node,
        ego_adapter,
        serial_interface,
        rviz_node
    ])
    
    return LaunchDescription([
        namespace_arg,
        use_sim_time_arg,
        drone_id_arg,
        map_size_x_arg,
        map_size_y_arg,
        map_size_z_arg,
        ego_3d_navigation
    ])